package info.yuehui.bigdatalimittest.util;

import lombok.extern.slf4j.Slf4j;
import org.apache.ibatis.session.ExecutorType;
import org.apache.ibatis.session.SqlSession;
import org.apache.ibatis.session.SqlSessionFactory;
import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
import org.springframework.stereotype.Component;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.util.CollectionUtils;

import javax.sql.DataSource;
import java.util.Collection;
import java.util.function.BiConsumer;

/**
 * @author zhangxing
 * @date 2020/03/15
 */
@Slf4j
@Component
public class DaoBatchUtils {

    /**
     * 默认一次提交1000条
     */
    public static final int DEFAULT_INSERT_STEP = 1000;

    private final SqlSessionFactory sqlSessionFactory;

    public DaoBatchUtils(SqlSessionFactory sqlSessionFactory) {
        this.sqlSessionFactory = sqlSessionFactory;
    }

    /**
     * 批量操作
     *
     * @param mapperClass mapper字节码类型
     * @param list        批量操作的数据
     * @param operation   操作
     * @param step        步长，即多少条sql一批次
     * @param <T>         批量操作的对象
     * @param <M>         批量操作使用的Mapper
     */
    @Transactional(rollbackFor = Throwable.class)
    public <T, M> void exec(Class<M> mapperClass, Collection<T> list, BiConsumer<M, T> operation, int step) {
        if (CollectionUtils.isEmpty(list)) {
            return;
        }

        // 批处理，并且不自动提交事务
        SqlSession sqlSession = sqlSessionFactory.openSession(ExecutorType.BATCH, false);
        M mapper = sqlSession.getMapper(mapperClass);

        int count = 0;
        for (T t : list) {
            count++;
            operation.accept(mapper, t);
            if (count % step == 0) {
                sqlSession.commit();
                sqlSession.clearCache();
                count = 0;
            }
        }

        if (count > 0) {
            sqlSession.commit();
            sqlSession.clearCache();
        }
    }

    /**
     * 批量操作，默认1000条一批
     *
     * @param mapperClass mapper字节码类型
     * @param list        批量操作的数据
     * @param operation   操作
     * @param <T>         批量操作的对象
     * @param <M>         批量操作使用的Mapper
     */
    @Transactional(rollbackFor = Throwable.class)
    public <T, M> void exec(Class<M> mapperClass, Collection<T> list, BiConsumer<M, T> operation) {
        exec(mapperClass, list, operation, DEFAULT_INSERT_STEP);
    }

}
